'''
根据本地数据，同步股票历史数据，保存到到本地文件, 共有3个参数: 

参数列表
    stock_code 股票代码, 格式为股市类型.股票编码, 如 'sh.603300'
    start_date 开始日期, 格式为 %Y-%m-%d, 如 '2024-01-01', 默认值为'2024-01-01'
    end_date 结束日期, 格式为 %Y-%m-%d, 如 '2024-12-31', 默认值为 None
    forced=False 是否强制写文件，如果保存的文件已存在，且forced为True时，会强制写入。

返回值
    
    无，文件会被保存在 f'./monthly_data/{stock_code}.csv', 如 ./monthly_data/sh.600749.csv
    
使用示例
    
    python S105.synchronize_monthly_data.py sh.600749 2020-01-01

    以上代码用于下载华铁应急近5年的月K线。

【参考资料】
1. 技术文档地址: https://ai-cyber-security.feishu.cn/docx/AB1fdqV9xo9K9axwT3hcMOOcnzc
2. BaoStock官网: http://baostock.com/baostock/index.php/Python_API文档
'''

import os, sys
import baostock as bs
from datetime import datetime, timedelta
import pandas as pd
from utils.common import *

def download_stock(lg, stock_code, sdate='2000-01-01', edate=None, out_dir='monthly_data', forced=False):
    # 读取数据, 默认开始日期为2000/01/01, 结束日期为当前日期
    column_names = "date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg"
    edate= datetime.now().strftime('%Y-%m-%d') if edate == None else edate
    data_list = []
    try:
        rs = bs.query_history_k_data_plus(stock_code, column_names, start_date=sdate, end_date=edate, frequency="m")
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
    except:
        pass
    return pd.DataFrame(data_list, columns=column_names.split(','))
# adjustflag：复权类型，默认不复权：3； 1：后复权；2：前复权。已支持分钟线、日线、周线、月线前后复权。 
#

############################### 主函数从这里开始 #############################################
print('功能说明: 根据已有数据同步下载指定股票或所有股票。')
print('         其中，默认数据文件夹为 ./daily_data\n')
print('函数调用:\n    python ./S104.synchonize_daily_data.py\n')

# 第1步 定义所有变量并初始化 -------------------------------
single_stock = ''                           # 支持对单一股票进行更新
stocknames_path = 'all_stock_codes.csv'     # 所有股票代码文件
data_dir = 'monthly_data'                     # 所有股票数据文件夹
code = 'SH.603538' if len(sys.argv) < 2 else sys.argv[1].upper()

if not os.path.exists(stocknames_path):
    print(f'股票所有代码文件 "{stocknames_path}" 不存在，请先调用 S101 下载所有股票代码.')
    sys.exit()    

if len(sys.argv) > 1:
    single_stock = sys.argv[1]

if not os.path.exists(data_dir):
    os.mkdir(data_dir)

begin_update_time = datetime.now()


# 第2步 用户登陆
lg = bs.login()
if  lg.error_code != '0':
    print("Loging failed. Error code: " + lg.error_msg)
    print("Loging failed. Error message: " + lg.error_code)
    sys.exit()
print(f'登陆成功，开始同步股票数据.')

#  第1步 数据处理
# 文件包括三列:  code, tradeStatus, code_name, 如：sz.301589 1 诺瓦星云
df_all_stock_codes = pd.read_csv(stocknames_path)
end_date = f'{datetime.now():%Y-%m-%d}'
finished = 0
total = len(df_all_stock_codes)
for index, row in df_all_stock_codes.iterrows():
    finished += 1
    result_str = 'downloaded' # 打印显示信息
    dr = row.tolist()   # print(index, dr) # 0 ['sh.000001', 1, '上证综合指数']
    stock_code = dr[0]
    stock_name = dr[2]
    stock_file = os.path.join(data_dir, f'{stock_code}.csv')
    start_date = '2010-01-01'
    column_names = "date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg".split(',')
    df1 = pd.DataFrame([], columns=column_names)
    try:
        if os.path.exists(stock_file):
            df1 = pd.read_csv(stock_file, parse_dates=['date'])
            last_date = (df1['date'].max() + timedelta(days=1)).strftime('%Y-%m-%d')
            start_date = last_date            
    except:
        pass
    
    
    df2 = download_stock(lg, dr[0], start_date, end_date, data_dir, forced=True)
    total_seconds = f'{(datetime.now()-begin_update_time).total_seconds():5.1f}'
    speed = finished/float(total_seconds) + 0.00000001  # rows/s
    left_seconds = (total - finished) / speed
    df3 = df1 if len(df2) == 0 else pd.concat([df1, df2], axis=0, ignore_index=True)
    df3.to_csv(stock_file, index=False)
    print(f'{index+1:4}/{total} {stock_code} 下载速度：{speed:5.1f} 条/秒，已用时：{total_seconds} 秒，剩余时间：{left_seconds:.0f} 秒, 已下载：{len(df2):4} 条 {stock_name}', flush=True)
    sys.stdout.flush()

try:
    bs.logout()
    print("登出成功")
except:
    print('推出失败')

print(f'所有股票数据同步完成，总更新时间：{(datetime.now() - begin_update_time).total_seconds():0.0}s.')
